from enum import Enum, unique
import numpy as np
from numpy.typing import NDArray
from typing import List, Dict, Optional, Union
import matplotlib.pyplot as plt
from abc import ABC, abstractmethod
from copy import deepcopy

from common_util import transform_3d_cloud
from common_util import point_in_polygon_roi, line_in_polygon_roi
from dataset_sdk import PointXYZ, BoundingBox
from common_util import MATPLOTLIB_COLOR, init_axes, matplot_text

__all__ = ['LaneDividerLineType', 'LaneDividerColor', 'LaneDivider',
           'LaneDividerJoint',
           'RoadEdgeCategory', 'RoadEdge',
           'RoadMarkingCategory', 'RoadMarking',
           'TrafficSignCategory', 'TrafficSign',
           'TrafficLightType', 'TrafficLightStatus', 'TrafficLight',
           'CenterLineTopoPoint', 'CenterLine', 'CenterLineConnect',
           'MapInfo']

def get_valid_axes(axes: Optional[plt.Axes]) -> plt.Axes:
    if axes is None:
        axes = init_axes()
        axes.axis('equal')
    return axes

class MapElement(ABC):
    @abstractmethod
    def __init__(self) -> None:
        pass

    @abstractmethod
    def to_dict(self) -> dict:
        pass

    @classmethod
    @abstractmethod
    def from_dict(cls, elem_dict: dict) -> Optional['MapElement']:
        pass

    @abstractmethod
    def shift_coordinate_system(self, translation: NDArray,
                                rotation_matrix: NDArray, inverse: bool=False) -> None:
        pass

    @abstractmethod
    def matplot(self, scale: float, rgb: str, alpha: float, axes: Optional[plt.Axes]=None) -> None:
        pass

    @abstractmethod
    def roi_filter(self, roi_points: NDArray) -> Optional['MapElement']:
        pass

class MapPointElement(MapElement):
    @abstractmethod
    def __init__(self, point: Union[tuple, list, NDArray]) -> None:
        assert len(point) == 3, 'ERROR: point should be (x, y, z)'
        self.point = PointXYZ(*point)

    def shift_coordinate_system(self, translation: NDArray,
                                rotation_matrix: NDArray, inverse: bool = False) -> None:
        self.point.shift_coordinate_system(translation, rotation_matrix, inverse)

    def matplot(self, scale: float, rgb: str, alpha: float, axes: Optional[plt.Axes]=None, text: str='') -> None:
        axes.plot(self.point.x, self.point.y, c=rgb, lw=scale, alpha=alpha)
        matplot_text(self.point.x, self.point.y, text, scale=2, rgb=rgb, alpha=alpha, axes=axes)

    def roi_filter(self, roi_points: NDArray) -> Optional['MapPointElement']:
        res = deepcopy(self)
        return res if point_in_polygon_roi(np.array(self.point.to_list()), roi_points) else None

class MapPointsElement(MapElement):
    @abstractmethod
    def __init__(self, points: NDArray) -> None:
        assert len(points.shape) == 2 and points.shape[-1] == 3, 'ERROR: points should be N*3'
        assert points.shape[0] > 1, 'ERROR: line object need at least 2 points'
        self.points = points

    def shift_coordinate_system(self, translation: NDArray,
                                rotation_matrix: NDArray, inverse: bool = False) -> None:
        self.points = transform_3d_cloud(self.points, translation, rotation_matrix, inverse)

    @abstractmethod
    def matplot(self, scale: float, rgb: str, alpha: float, axes: Optional[plt.Axes]=None, text: str='') -> None:
        axes.plot(self.points[:, 0], self.points[:, 1], c=rgb, lw=scale, alpha=alpha)
        matplot_text(self.points[0, 0], self.points[0, 1], text, scale=2, rgb=rgb, alpha=alpha, axes=axes)

    def roi_filter(self, roi_points: NDArray) -> Optional['MapPointsElement']:
        res = deepcopy(self)
        res.points = line_in_polygon_roi(self.points, roi_points)[0]
        return res if len(res.points) > 1 else None

class MapPolygonElement(MapElement):
    @abstractmethod
    def __init__(self, polygon: NDArray) -> None:
        assert len(polygon.shape) == 2 and polygon.shape[-1] == 3, 'ERROR: polygon should be N*3'
        self.polygon = polygon

    def shift_coordinate_system(self, translation: NDArray,
                                rotation_matrix: NDArray, inverse: bool = False) -> None:
        self.polygon = transform_3d_cloud(self.polygon, translation, rotation_matrix, inverse)

    @abstractmethod
    def matplot(self, scale: float, rgb: str, alpha: float, axes: Optional[plt.Axes]=None, text: str='') -> None:
        axes.fill(self.polygon[:, 0], self.polygon[:, 1], c=rgb, lw=scale, alpha=alpha)
        matplot_text(self.polygon[0, 0], self.polygon[0, 1], text, scale=2, rgb=rgb, alpha=alpha, axes=axes)

    def roi_filter(self, roi_points: NDArray) -> Optional['MapPolygonElement']:
        res = deepcopy(self)
        points_in_roi = line_in_polygon_roi(self.polygon, roi_points)[0]
        return res if len(points_in_roi) > 1 else None

class MapBoxElement(MapElement):
    @abstractmethod
    def __init__(self, center, size, rotation) -> None:
        assert len(center) == len(size) == len(rotation) == 3, 'ERROR: box info format wrong'
        self.box = BoundingBox(center, size, rotation)

    def shift_coordinate_system(self, translation: NDArray,
                                rotation_matrix: NDArray, inverse: bool = False) -> None:
        self.box.shift_coordinate_system(translation, rotation_matrix, inverse)

    @abstractmethod
    def matplot(self, scale: float, rgb: str, alpha: float, axes: Optional[plt.Axes]=None, text: str='') -> None:
        self.box.matplot_on_bev(scale, rgb, alpha, axes)
        matplot_text(self.box.center.x, self.box.center.y, text, scale=2, rgb=rgb, alpha=alpha, axes=axes)

    def roi_filter(self, roi_points: NDArray) -> Optional['MapBoxElement']:
        res = deepcopy(self)
        return res if point_in_polygon_roi(np.array(self.box.center.to_list()), roi_points) else None

@unique
class LaneDividerLineType(Enum):
    solid = 'solid'
    dashed = 'dashed'
    unknown = 'unknown'

@unique
class LaneDividerColor(Enum):
    white = 'white'
    yellow = 'yellow'
    blue = 'blue'
    green = 'green'
    red = 'red'
    unknown = 'unknown'

class LaneDivider(MapPointsElement):
    def __init__(self, id_, line_type: LaneDividerLineType, color: LaneDividerColor, points: NDArray) -> None:
        assert isinstance(line_type, LaneDividerLineType)
        assert isinstance(color, LaneDividerColor)
        self.id = id_
        self.line_type: LaneDividerLineType = line_type
        self.color: LaneDividerColor = color
        super().__init__(points)

    def to_dict(self) -> dict:
        return {
            'id': self.id,
            'line_type': self.line_type.value,
            'color': self.color.value,
            'points': self.points.tolist(),
        }

    @classmethod
    def from_dict(cls, elem_dict: dict) -> Optional['LaneDivider']:
        if len(elem_dict['points']) == 0:
            return None
        return cls(elem_dict['id'], LaneDividerLineType(elem_dict['line_type']),
                   LaneDividerColor(elem_dict['color']), np.array(elem_dict['points']))

    def matplot(self, scale=0.2, rgb=MATPLOTLIB_COLOR['yellow'], alpha=0.5, axes: Optional[plt.Axes]=None) -> None:
        axes = get_valid_axes(axes)
        super().matplot(scale, rgb, alpha, axes, f'{self.line_type.value}\n{self.color.value}')

class LaneDividerJoint(MapPointElement):
    def __init__(self, lane_divider_id1, lane_divider_id2, point: Union[tuple, list, NDArray]) -> None:
        self.lane_divider_id1 = lane_divider_id1
        self.lane_divider_id2 = lane_divider_id2
        super().__init__(point)

    def to_dict(self) -> dict:
        return {
            'lane_divider_id1': self.lane_divider_id1,
            'lane_divider_id2': self.lane_divider_id2,
            'point': self.point.to_list(),
        }

    @classmethod
    def from_dict(cls, elem_dict: dict) -> Optional['LaneDividerJoint']:
        return cls(elem_dict['lane_divider_id1'], elem_dict['lane_divider_id2'], elem_dict['point'])

    def matplot(self, scale=0.3, rgb=MATPLOTLIB_COLOR['orange'], alpha=0.5, axes: Optional[plt.Axes]=None) -> None:
        axes = get_valid_axes(axes)
        super().matplot(scale, rgb, alpha, axes)

@unique
class RoadEdgeCategory(Enum):
    road_edge = 'road_edge'
    traffic_barrier = 'traffic_barrier'
    unknown = 'unknown'

class RoadEdge(MapPointsElement):
    def __init__(self, id_, category: RoadEdgeCategory, points: NDArray) -> None:
        self.id = id_
        self.category: RoadEdgeCategory = category
        super().__init__(points)

    def to_dict(self) -> dict:
        return {
            'id': self.id,
            'category': self.category.value,
            'points': self.points.tolist(),
        }

    @classmethod
    def from_dict(cls, elem_dict: dict) -> Optional['RoadEdge']:
        if len(elem_dict['points']) == 0:
            return None
        return cls(elem_dict['id'], RoadEdgeCategory(elem_dict['category']), np.array(elem_dict['points']))

    def matplot(self, scale=0.2, rgb=MATPLOTLIB_COLOR['black'], alpha=0.5, axes: Optional[plt.Axes]=None) -> None:
        axes = get_valid_axes(axes)
        super().matplot(scale, rgb, alpha, axes, self.category.value)

@unique
class RoadMarkingCategory(Enum):
    straight = 'straight'
    left = 'left'
    right = 'right'
    u_turn = 'u_turn'
    cross = 'cross'
    no_stop = 'no_stop'
    crosswalk = 'crosswalk'
    stop_line = 'stop_line'
    speed_bump = 'speed_bump'
    confluence_left = 'confluence_left'
    confluence_right = 'confluence_right'

class RoadMarking(MapPolygonElement):
    def __init__(self, id_, category: List[RoadMarkingCategory], polygon: NDArray) -> None:
        for i in category:
            assert isinstance(i, RoadMarkingCategory)
        self.id = id_
        self.category: List[RoadMarkingCategory] = category
        super().__init__(polygon)

    def to_dict(self) -> dict:
        return {
            'id': self.id,
            'category': [i.value for i in self.category],
            'polygon': self.polygon.tolist(),
        }

    @classmethod
    def from_dict(cls, elem_dict: dict) -> Optional['RoadMarking']:
        if len(elem_dict['polygon']) == 0:
            return None
        category = [RoadMarkingCategory(it) for it in elem_dict['category']]
        return cls(elem_dict['id'], category, np.array(elem_dict['polygon']))

    def matplot(self, scale=0.2, rgb=MATPLOTLIB_COLOR['pink'], alpha=0.5, axes: Optional[plt.Axes]=None) -> None:
        axes = get_valid_axes(axes)
        super().matplot(scale, rgb, alpha, axes, '_'.join([i.value for i in self.category]))

@unique
class TrafficSignCategory(Enum):
    no_entry = 'no_entry'
    no_vehicle = 'no_vehicle'
    no_stop = 'no_stop'
    stop = 'stop'
    speed_unlimit = 'speed_unlimit'
    speed_limit_5 = 'speed_limit_5'
    speed_limit_10 = 'speed_limit_10'
    speed_limit_15 = 'speed_limit_15'
    speed_limit_20 = 'speed_limit_20'
    speed_limit_30 = 'speed_limit_30'
    speed_limit_40 = 'speed_limit_40'
    speed_limit_50 = 'speed_limit_50'
    speed_limit_60 = 'speed_limit_60'
    speed_limit_70 = 'speed_limit_70'
    speed_limit_80 = 'speed_limit_80'
    speed_limit_90 = 'speed_limit_90'
    speed_limit_100 = 'speed_limit_100'
    speed_limit_110 = 'speed_limit_110'
    speed_limit_120 = 'speed_limit_120'
    speed_limit_130 = 'speed_limit_130'
    lowest_speed_50 = 'lowest_speed_50'
    lowest_speed_60 = 'lowest_speed_60'
    lowest_speed_70 = 'lowest_speed_70'
    lowest_speed_80 = 'lowest_speed_80'
    lowest_speed_90 = 'lowest_speed_90'
    lowest_speed_100 = 'lowest_speed_100'
    lowest_speed_110 = 'lowest_speed_110'
    no_overtaking = 'no_overtaking'
    end_of_no_overtaking = 'end_of_no_overtaking'
    pole = 'pole'
    sign = 'sign'

class TrafficSign(MapBoxElement):
    def __init__(self, id_, category: TrafficSignCategory, center, size, rotation) -> None:
        assert isinstance(category, TrafficSignCategory)
        self.id = id_
        self.category: TrafficSignCategory = category
        super().__init__(center, size, rotation)

    def to_dict(self) -> dict:
        self_dict = {
            'id': self.id,
            'category': self.category.value
        }
        self_dict.update(self.box.to_dict())
        return self_dict

    @classmethod
    def from_dict(cls, elem_dict: dict) -> Optional['TrafficSign']:
        return cls(elem_dict['id'], TrafficSignCategory(elem_dict['category']),
                   [elem_dict['center']['x'], elem_dict['center']['y'], elem_dict['center']['z']],
                   [elem_dict['size']['x'], elem_dict['size']['y'], elem_dict['size']['z']],
                   [elem_dict['rotation']['x'], elem_dict['rotation']['y'], elem_dict['rotation']['z']])

    def matplot(self, scale=0.2, rgb=MATPLOTLIB_COLOR['lightred'], alpha=0.5, axes: Optional[plt.Axes]=None) -> None:
        axes = get_valid_axes(axes)
        super().matplot(scale, rgb, alpha, axes, f'{self.category.value}')

@unique
class TrafficLightType(Enum):
    round_ = 'round'
    straight = 'straight'
    left = 'left'
    right = 'right'
    u_turn = 'u_turn'
    straight_left = 'straight_left'
    straight_right = 'straight_right'
    left_u_turn = 'left_u_turn'
    access = 'access'
    unknown = 'unknown'

@unique
class TrafficLightStatus(Enum):
    none = 'none'
    red = 'red'
    yellow = 'yellow'
    green = 'green'
    unknown = 'unknown'

class TrafficLight(MapBoxElement):
    def __init__(self, id_, type_: TrafficLightType, status: TrafficLightStatus, center, size, rotation) -> None:
        assert isinstance(type_, TrafficLightType)
        assert isinstance(status, TrafficLightStatus)
        self.id = id_
        self.type: TrafficLightType = type_
        self.status: TrafficLightStatus = status
        super().__init__(center, size, rotation)

    def to_dict(self) -> dict:
        self_dict = {
            'id': self.id,
            'type': self.type.value,
            'status': self.status.value,
        }
        self_dict.update(self.box.to_dict())
        return self_dict

    @classmethod
    def from_dict(cls, elem_dict: dict) -> Optional['TrafficLight']:
        return cls(elem_dict['id'], TrafficLightType(elem_dict['type']), TrafficLightStatus(elem_dict['status']),
                   [elem_dict['center']['x'], elem_dict['center']['y'], elem_dict['center']['z']],
                   [elem_dict['size']['x'], elem_dict['size']['y'], elem_dict['size']['z']],
                   [elem_dict['rotation']['x'], elem_dict['rotation']['y'], elem_dict['rotation']['z']])

    def matplot(self, scale=0.2, rgb=MATPLOTLIB_COLOR['cyan'], alpha=0.5, axes: Optional[plt.Axes]=None) -> None:
        axes = get_valid_axes(axes)
        super().matplot(scale, rgb, alpha, axes, f'{self.type.value}\n{self.status.value}')

class CenterLineTopoPoint(MapPointElement):
    def __init__(self, id_, center_line_id, index: int, point: Union[tuple, list, NDArray]) -> None:
        self.id = id_
        self.center_line_id = center_line_id
        self.index = index
        super().__init__(point)

    def to_dict(self) -> dict:
        return {
            'id': self.id,
            'center_line_id': self.center_line_id,
            'index': self.index,
            'point': self.point.to_list(),
        }

    @classmethod
    def from_dict(cls, elem_dict: dict) -> Optional['CenterLineTopoPoint']:
        return cls(elem_dict['id'], elem_dict['center_line_id'], elem_dict['index'], elem_dict['point'])

    def matplot(self, scale=0.2, rgb=MATPLOTLIB_COLOR['blue'], alpha=0.5, axes: Optional[plt.Axes]=None) -> None:
        axes = get_valid_axes(axes)
        super().matplot(scale, rgb, alpha, axes, self.id)

class CenterLine(MapPointsElement):
    def __init__(self, id_, topo_point_ids: List[str], points: NDArray) -> None:
        self.id = id_
        self.topo_point_ids = topo_point_ids
        super().__init__(points)

    def to_dict(self) -> dict:
        return {
            'id': self.id,
            'topo_point_ids': self.topo_point_ids,
            'points': self.points.tolist(),
        }

    @classmethod
    def from_dict(cls, elem_dict: dict) -> Optional['CenterLine']:
        if len(elem_dict['points']) == 0:
            return None
        return cls(elem_dict['id'], elem_dict['topo_point_ids'], np.array(elem_dict['points']))

    def matplot(self, scale=0.2, rgb=MATPLOTLIB_COLOR['red'], alpha=0.5, axes: Optional[plt.Axes]=None) -> None:
        axes = get_valid_axes(axes)
        super().matplot(scale, rgb, alpha, axes, str(self.id))

class CenterLineConnect(MapPointsElement):
    def __init__(self, id_, pred_line, succ_line, points: NDArray) -> None:
        self.id = id_
        self.pred_line = pred_line
        self.succ_line = succ_line
        super().__init__(points)

    def to_dict(self) -> dict:
        return {
            'id': self.id,
            'pred_line': self.pred_line,
            'succ_line': self.succ_line,
            'points': self.points.tolist(),
        }

    @classmethod
    def from_dict(cls, elem_dict: dict) -> Optional['CenterLineConnect']:
        return cls(elem_dict['id'], elem_dict['pred_line'], elem_dict['succ_line'], np.array(elem_dict['points']))

    def matplot(self, scale=0.2, rgb=MATPLOTLIB_COLOR['lightblue'], alpha=0.5, axes: Optional[plt.Axes]=None) -> None:
        axes = get_valid_axes(axes)
        super().matplot(scale, rgb, alpha, axes, str(self.id))

MAP_ATTRIBUTE_CLASS: Dict[str, MapElement] = {
    'lane_dividers': LaneDivider,
    'lane_divider_joints': LaneDividerJoint,
    'road_edges': RoadEdge,
    'road_markings': RoadMarking,
    'traffic_signs': TrafficSign,
    'traffic_lights': TrafficLight,
    'center_line_topo_points': CenterLineTopoPoint,
    'center_lines': CenterLine,
    'center_line_connects': CenterLineConnect,
}

class MapInfo():
    def __init__(self) -> None:
        self.lane_dividers: List[LaneDivider] = []
        self.lane_divider_joints: List[LaneDividerJoint] = []
        self.road_edges: List[RoadEdge] = []
        self.road_markings: List[RoadMarking] = []
        self.traffic_signs: List[TrafficSign] = []
        self.traffic_lights: List[TrafficLight] = []
        self.center_line_topo_points: List[CenterLineTopoPoint] = []
        self.center_lines: List[CenterLine] = []
        self.center_line_connects: List[CenterLineConnect] = []

    def to_dict(self) -> dict:
        res = dict()
        for name, value in vars(self).items():
            value: List[MapElement] = value
            res[name] = [i.to_dict() for i in value]
        return res

    @classmethod
    def from_dict(cls, map_dict: Dict[str, List[dict]]) -> 'MapInfo':
        map_info = cls()
        for attr, elem_list in map_dict.items():
            map_attr: List[MapElement] = getattr(map_info, attr)
            for elem_dict in elem_list:
                elem = MAP_ATTRIBUTE_CLASS[attr].from_dict(elem_dict)
                if elem is not None:
                    map_attr.append(elem)
        return map_info

    def shift_coordinate_system(self, translation: NDArray,
                                rotation_matrix: NDArray, inverse: bool=False) -> None:
        for value in vars(self).values():
            value: List[MapElement] = value
            for i in value:
                i.shift_coordinate_system(translation, rotation_matrix, inverse)

    def roi_filter(self, roi_points: NDArray) -> 'MapInfo':
        res = MapInfo()
        for attr, value in vars(self).items():
            res_attr: List[MapElement] = getattr(res, attr)
            value: List[MapElement] = value
            for item in value:
                roi_item = item.roi_filter(roi_points)
                if roi_item is not None:
                    res_attr.append(roi_item)
        return res

    @staticmethod
    def generate_roi_points(roi: Union[float, tuple]) -> NDArray:
        if isinstance(roi, (int, float)):
            from math import cos, sin, pi
            ROI_SAMPLE = 32
            return np.array([[cos(pi/ROI_SAMPLE*2*i)*roi, sin(pi/ROI_SAMPLE*2*i)*roi] for i in range(ROI_SAMPLE)])
        elif len(roi) == 2:
            return np.array([[-roi[0], -roi[1]], [-roi[0], roi[1]], [roi[0], roi[1]], [roi[0], -roi[1]]])
        elif len(roi) == 4:
            return np.array([[roi[0], roi[1]], [roi[2], roi[1]], [roi[2], roi[3]], [roi[0], roi[3]]])
        else:
            raise NotImplementedError(f'ERROR: roi should be 1, 2 or 4 float')
